import math
from Env.action import *

class UnitStatus:
    def __init__(self, status, idx):
        self.status = status
        self.idx = idx
#

# 动作执行
class ActionGenerator:
    def __init__(self):
        print("Creat ActionGenerator Instance")
        self.state = None
        self.action_function_map = {
            "":                         self.NoneFunc,
        }

    def generate(self, status_list):
        self.cmds = []
        for status in status_list:
            cmd = self.action_function_map[status.status](status.idx)
            if cmd != "":
                self.cmds.append(cmd)
        return self.cmds

    def NoneFunc(self, idx):
        return ''
        

class StatusTransfer:
    def __init__(self):
        print("Creat StatusTransfer Instance")
        self.trans_function_map = {
            "UavInit":                  self.uav_init,
        }
        self.reset()
    
    def reset(self):
        self.state = None

    def transfer(self, status_list):
        if self.cross_river_check:
            self.CrossRiverCheck()
        new_status_list = []
        for status in status_list:
            new_status = self.trans_function_map[status.status](status.idx)
            if new_status:
                new_status_list.append(UnitStatus(new_status, status.idx))
        return new_status_list

  
  
class FSMController:
    def __init__(self):
        self.action_executor = ActionGenerator()
        self.status_transfer = StatusTransfer()
        print("Creat a Controller Instance")

    def init_status(self):
        status = []
        [status.append(UnitStatus("UavInit", i)) for i in range(3)]
        [status.append(UnitStatus("TankInit", i)) for i in range(4)]
        [status.append(UnitStatus("SoldInit", i)) for i in range(6)]
        self.status = status

    def reset(self):
        self.init_status()

    def step(self, state):
        self.status_transfer.state = state
        self.action_executor.state = state
        print('------------- step ----------------')
        for s in self.status:
            print(s.idx, s.status)
        #print(self.status)
        self.status = self.status_transfer.transfer(self.status)
        cmds = self.action_executor.generate(self.status)
        return cmds


def main():
    print("Exec Main Function")
    import os
    import time
    from Env.client import EnvClient
    client = EnvClient(0,0,0,0)
    client.reset()
    obs = client.get_observation()
    player = FSMController()
    player.reset()
    while True:
        cmds = player.step(obs)
        print(cmds)
        client.take_action(cmds)
        time.sleep(2)
        obs = client.get_observation()
        #obs.print_state()


if __name__ == '__main__':
    main()
